#pragma once

#include <cuda_runtime.h>
#include <iostream>

#define CUDA_CHECK(call)                                          \
  do {                                                            \
    cudaError_t error = call;                                     \
    if (error != cudaSuccess) {                                   \
      std::cerr << "CUDA error: " << cudaGetErrorString(error)    \
                << " at " << __FILE__ << ":" << __LINE__ << std::endl; \
      exit(1);                                                    \
    }                                                             \
  } while (0)

inline int getOptimalThreadsPerBlock() {
    int device;
    CUDA_CHECK(cudaGetDevice(&device));
    cudaDeviceProp prop;
    CUDA_CHECK(cudaGetDeviceProperties(&prop, device));
    int maxThreadsPerBlock = prop.maxThreadsPerBlock;
    int threadsPerBlock = prop.warpSize;
    
    // 可以根据需要添加其他逻辑来调整 threadsPerBlock
    // 例如，基于计算能力 (prop.major, prop.minor)
    // 或共享内存大小 (prop.sharedMemPerBlock)

    if (maxThreadsPerBlock >= 1024) {
      threadsPerBlock = 1024; 
    } else if (maxThreadsPerBlock >= 512) {
        threadsPerBlock = 512; 
    } else if (maxThreadsPerBlock >= 256) {
        threadsPerBlock = 256; 
    } else if (maxThreadsPerBlock >= 128) {
      threadsPerBlock = 128; 
    }
    return threadsPerBlock;
}